import os
import sys
import requests
from urllib.parse import urljoin
from concurrent.futures import ThreadPoolExecutor
from Crypto.Cipher import AES
import shutil
import re
import binascii
import logging
import subprocess
import time
import tkinter as tk
from tkinter import ttk, filedialog, messagebox, scrolledtext
import threading

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

class M3U8DownloaderApp:
    def __init__(self, root):
        self.root = root
        self.root.title("M3U8视频下载器")
        self.root.geometry("800x600")
        
        self.create_widgets()
        self.stop_event = threading.Event()
        self.running = False

    def create_widgets(self):
        main_frame = ttk.Frame(self.root, padding=10)
        main_frame.pack(fill=tk.BOTH, expand=True)

        input_frame = ttk.LabelFrame(main_frame, text="下载设置", padding=10)
        input_frame.pack(fill=tk.X, pady=5)

        # URL输入
        ttk.Label(input_frame, text="M3U8 URL:").grid(row=0, column=0, sticky=tk.W)
        self.url_entry = ttk.Entry(input_frame, width=60)
        self.url_entry.grid(row=0, column=1, padx=5)
        ttk.Button(input_frame, text="粘贴", command=self.paste_url).grid(row=0, column=2, padx=5)

        # TS临时路径
        ttk.Label(input_frame, text="TS存储路径:").grid(row=1, column=0, sticky=tk.W)
        self.ts_path_entry = ttk.Entry(input_frame, width=60)
        self.ts_path_entry.grid(row=1, column=1, padx=5)
        ttk.Button(input_frame, text="浏览", command=self.browse_ts_path).grid(row=1, column=2, padx=5)

        # 输出目录
        ttk.Label(input_frame, text="输出目录:").grid(row=2, column=0, sticky=tk.W)
        self.output_dir_entry = ttk.Entry(input_frame, width=60)
        self.output_dir_entry.grid(row=2, column=1, padx=5)
        ttk.Button(input_frame, text="浏览", command=self.browse_output_dir).grid(row=2, column=2, padx=5)

        # 文件名
        ttk.Label(input_frame, text="输出文件名:").grid(row=3, column=0, sticky=tk.W)
        self.ts_filename_entry = ttk.Entry(input_frame, width=60)
        self.ts_filename_entry.insert(0, time.strftime("%Y-%m-%d"))
        self.ts_filename_entry.grid(row=3, column=1, padx=5)

        # 控制按钮
        btn_frame = ttk.Frame(main_frame)
        btn_frame.pack(pady=10)
        self.start_btn = ttk.Button(btn_frame, text="开始下载", command=self.start_download)
        self.start_btn.pack(side=tk.LEFT, padx=5)
        self.stop_btn = ttk.Button(btn_frame, text="停止", command=self.stop_download, state=tk.DISABLED)
        self.stop_btn.pack(side=tk.LEFT, padx=5)

        # 进度条
        self.progress = ttk.Progressbar(main_frame, orient=tk.HORIZONTAL, mode='determinate')
        self.progress.pack(fill=tk.X, pady=5)

        # 日志框
        log_frame = ttk.LabelFrame(main_frame, text="日志", padding=10)
        log_frame.pack(fill=tk.BOTH, expand=True)
        self.log_text = scrolledtext.ScrolledText(log_frame, height=10)
        self.log_text.pack(fill=tk.BOTH, expand=True)
        self.log_handler = TextHandler(self.log_text)
        logging.getLogger().addHandler(self.log_handler)

    def paste_url(self):
        self.url_entry.delete(0, tk.END)
        self.url_entry.insert(0, self.root.clipboard_get())

    def browse_ts_path(self):
        path = filedialog.askdirectory()
        if path:
            self.ts_path_entry.delete(0, tk.END)
            self.ts_path_entry.insert(0, path)

    def browse_output_dir(self):
        path = filedialog.askdirectory()
        if path:
            self.output_dir_entry.delete(0, tk.END)
            self.output_dir_entry.insert(0, path)

    def start_download(self):
        if self.running:
            return
            
        m3u8_url = self.url_entry.get()
        ts_folder = self.ts_path_entry.get() or time.strftime("%Y-%m-%d")
        output_dir = self.output_dir_entry.get() or os.getcwd()
        ts_filename = self.ts_filename_entry.get() or time.strftime("%Y-%m-%d")
        
        if not m3u8_url:
            messagebox.showerror("错误", "请输入M3U8 URL")
            return
            
        self.running = True
        self.stop_event.clear()
        self.start_btn.config(state=tk.DISABLED)
        self.stop_btn.config(state=tk.NORMAL)
        self.progress["value"] = 0
        
        threading.Thread(target=self.run_download, args=(m3u8_url, ts_folder, output_dir, ts_filename)).start()

    def stop_download(self):
        self.stop_event.set()
        self.running = False
        self.start_btn.config(state=tk.NORMAL)
        self.stop_btn.config(state=tk.DISABLED)
        logging.info("用户请求停止下载")

    def run_download(self, m3u8_url, ts_folder, output_dir, ts_filename):
        try:
            base_url = os.path.dirname(m3u8_url) + '/'
            m3u8_content = self.download_m3u8(m3u8_url)
            ts_urls, key, iv = self.parse_m3u8(m3u8_content, base_url)
            
            self.progress["maximum"] = len(ts_urls)
            
            with ThreadPoolExecutor(max_workers=5) as executor:
                futures = []
                for index, ts_url in enumerate(ts_urls):
                    if self.stop_event.is_set():
                        break
                    futures.append(executor.submit(
                        self.download_ts, ts_url, index, key, iv, ts_folder
                    ))
                
                for i, future in enumerate(futures):
                    try:
                        future.result()
                        self.progress["value"] = i + 1
                        self.root.update_idletasks()
                    except Exception as e:
                        logging.error(f"下载失败: {e}")

            if not self.stop_event.is_set():
                output_filename = os.path.join(output_dir, ts_filename)
                self.merge_ts_files(len(ts_urls), ts_folder, output_filename + ".ts")
                self.convert_to_mp4(output_filename + ".ts", output_filename + ".mp4")
                
        except Exception as e:
            logging.error(f"发生错误: {str(e)}")
            messagebox.showerror("错误", f"操作失败: {str(e)}")
        finally:
            try:
                self.cleanup_files(ts_folder)
            except Exception as cleanup_error:
                logging.error(f"最终清理失败: {str(cleanup_error)}")
            self.running = False
            self.start_btn.config(state=tk.NORMAL)
            self.stop_btn.config(state=tk.DISABLED)

    def download_m3u8(self, m3u8_url):
        try:
            response = requests.get(m3u8_url, timeout=10)
            response.raise_for_status()
            m3u8_content = response.text
            with open('playlist.m3u8', 'w') as f:
                f.write(m3u8_content)
            logging.info("已下载m3u8文件")
            return m3u8_content
        except Exception as e:
            logging.error(f"下载m3u8文件失败: {e}")
            raise

    def parse_m3u8(self, m3u8_content, base_url):
        ts_urls = []
        key = None
        iv = None
        lines = m3u8_content.strip().split('\n')
        for index, line in enumerate(lines):
            if line.startswith("#EXT-X-KEY"):
                method = re.search(r'METHOD=([A-Za-z0-9-]+)', line)
                uri = re.search(r'URI="([^"]+)"', line)
                iv_match = re.search(r'IV=0x([0-9a-fA-F]+)', line)

                if method and method.group(1) != 'AES-128':
                    raise ValueError(f"不支持的加密方法: {method.group(1)}")
                if uri:
                    key_url = uri.group(1)
                    key_url = urljoin(base_url, key_url)
                    key = self.download_key(key_url)
                if iv_match:
                    iv_hex = iv_match.group(1)
                    iv = binascii.unhexlify(iv_hex)
                else:
                    iv = None
            elif line and not line.startswith("#"):
                ts_url = urljoin(base_url, line)
                ts_urls.append(ts_url)
        return ts_urls, key, iv

    def download_key(self, key_url):
        try:
            response = requests.get(key_url, timeout=10)
            response.raise_for_status()
            with open('key.key', 'wb') as f:
                f.write(response.content)
            logging.info("已下载并保存key")
            return response.content
        except Exception as e:
            logging.error(f"下载key失败: {e}")
            raise

    def download_ts(self, url, index, key, iv, ts_folder):
        try:
            response = requests.get(url, timeout=10)
            response.raise_for_status()
            ts_content = response.content
            
            if key:
                if iv is None:
                    iv = index.to_bytes(16, byteorder='big')
                cipher = AES.new(key, AES.MODE_CBC, iv=iv)
                ts_content = cipher.decrypt(ts_content)
                
            os.makedirs(ts_folder, exist_ok=True)
            ts_filename = os.path.join(ts_folder, f"{index}.ts")
            with open(ts_filename, 'wb') as f:
                f.write(ts_content)
            logging.info(f"已下载并保存 {ts_filename}")
        except Exception as e:
            logging.error(f"下载/解密TS文件失败: {e}")
            raise

    def merge_ts_files(self, num_files, ts_folder, output_path):
        try:
            with open(output_path, 'wb') as merged:
                for i in range(num_files):
                    ts_path = os.path.join(ts_folder, f"{i}.ts")
                    if os.path.exists(ts_path):
                        with open(ts_path, 'rb') as f:
                            shutil.copyfileobj(f, merged)
                        logging.debug(f"已合并 {ts_path}")
                    else:
                        logging.warning(f"文件 {ts_path} 不存在，跳过")
            logging.info(f"合并完成: {output_path}")
        except Exception as e:
            logging.error(f"合并失败: {e}")
            raise

    def convert_to_mp4(self, input_path, output_path):
        try:
            subprocess.run(
                ['ffmpeg', '-y', '-i', input_path, 
                 '-c:v', 'copy', '-c:a', 'copy', output_path],
                check=True,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL
            )
            os.remove(input_path)
            logging.info(f"转换完成: {output_path}")
        except Exception as e:
            logging.error(f"转换失败: {e}")
            raise

    def cleanup_files(self, ts_folder):
        """清理所有临时文件"""
        try:
            # 清理TS文件
            if os.path.exists(ts_folder):
                ts_count = 0
                for filename in os.listdir(ts_folder):
                    if filename.lower().endswith(".ts"):
                        file_path = os.path.join(ts_folder, filename)
                        try:
                            os.remove(file_path)
                            ts_count += 1
                        except Exception as e:
                            logging.error(f"删除失败 {file_path}: {str(e)}")
                logging.info(f"已清理 {ts_count} 个TS文件")

            # 清理其他临时文件
            temp_files = ['key.key', 'playlist.m3u8']
            for temp_file in temp_files:
                if os.path.exists(temp_file):
                    try:
                        os.remove(temp_file)
                        logging.info(f"已删除临时文件: {temp_file}")
                    except Exception as e:
                        logging.error(f"删除失败 {temp_file}: {str(e)}")
        except Exception as e:
            logging.error(f"清理过程中发生错误: {str(e)}")

class TextHandler(logging.Handler):
    def __init__(self, text_widget):
        super().__init__()
        self.text_widget = text_widget
        
    def emit(self, record):
        msg = self.format(record)
        self.text_widget.configure(state="normal")
        self.text_widget.insert(tk.END, msg + "\n")
        self.text_widget.configure(state="disabled")
        self.text_widget.yview(tk.END)

if __name__ == "__main__":
    root = tk.Tk()
    app = M3U8DownloaderApp(root)
    root.mainloop()